import pandas as pd
import numpy as np
import math



def load_tcpdump_pcap_csv(filename, length = None):
    if length == None:
        df_data = pd.read_csv(filename)
    else:
        df_data = pd.read_csv(filename)[0:length]
    return df_data


def get_packets_in_flight(filename, SRC, DST, length = None):
    df_data = load_tcpdump_pcap_csv(filename, length)
    inflight_packets = 0
    inflight_packets_list = []
    current_time_list = []
    last_ack = 0

    last_seq = 0

    for index, row in df_data.iterrows():
        time_relative = row['frame.time_relative']
        current_time_list.append(time_relative)
        ack = row['tcp.ack']
        tcp_len = row['tcp.len']
        seq = row['tcp.seq']
        src = row['ip.src']
        dst = row['ip.dst']
    
        if dst == DST and ack:
        # if ack == 1:
            # inflight_packets += tcp_len
            inflight_packets += seq - last_seq
            last_seq = seq
        elif src == DST and ack:
            if ack and ack - last_ack >= 0:
                inflight_packets -= ack - last_ack
                last_ack = ack
            else:
                print('tmp debug', ack - last_ack, ack, last_ack)
        
        # Bytes -> MBytes
        inflight_packets_list.append(inflight_packets / 1000000)
    
    return current_time_list, inflight_packets_list

import csv
fields = ["frame.number","_ws.col.Time","frame.len","tcp.seq","ip.src","tcp.srcport","ip.dst","tcp.dstport","tcp.len","tcp.ack","frame.time_relative"]
class pkt:
    contents=[]
    def __init__(self, fields) -> None:
        self.contents=[]
        for f in fields:
            self.contents.append(f)

    def get(self, field):
        return self.contents[fields.index(field)]
    
# From Ayush    
def process_flows(filename, SRC, DST):
    with open(filename) as csv_file:
        csv_reader = csv.reader(csv_file, delimiter=',')
        line_count = 0
        total_bytes=0
        '''
        Flow tracking:
        o Identify all packets that are either sourced from or headed to 100.64.0.2
        o Group different flows by client's port
        '''
        flows={}
        for row in csv_reader:
            packet=pkt(row)

            validPkt=False
            if line_count==0:
                # reject the header
                line_count+=1
                continue
            if packet.get("ip.src")==DST and packet.get("frame.time_relative")!='' and packet.get("tcp.ack")!='': 
                # we care about this ACK packet
                validPkt=True
                port=packet.get("tcp.srcport")
                
                if port not in flows:
                    flows[port]={"serverip":packet.get("ip.dst"), "serverport":packet.get("tcp.destport"), "times":[], "windows":[], "cwnd":[], "bif":0, "last_ack":0, "last_seq":0, "pif":0}
                flows[port]["times"].append(float(packet.get("frame.time_relative")))
                flows[port]["bif"]-=(int(packet.get("tcp.ack"))-int(flows[port]["last_ack"]))
                flows[port]["last_ack"]=int(packet.get("tcp.ack"))
                flows[port]["pif"]-=1
                flows[port]["cwnd"].append(flows[port]["pif"])
                flows[port]["windows"].append(int(flows[port]["bif"]))

            elif packet.get("ip.dst")==DST and packet.get("frame.time_relative")!='' and packet.get("tcp.seq")!='':
                #we care about this Data packet
                validPkt=True
                port=packet.get("tcp.dstport")
                if port not in flows:
                    flows[port]={"serverip":packet.get("ip.src"), "serverport":packet.get("tcp.srcport"), "times":[], "windows":[], "cwnd":[], "bif":0, "last_ack":0, "last_seq":0, "pif":0}
                flows[port]["times"].append(float(packet.get("frame.time_relative")))
                flows[port]["bif"]+=(int(packet.get("tcp.seq"))-int(flows[port]["last_seq"]))
                flows[port]["last_seq"]=int(packet.get("tcp.seq"))
                flows[port]["pif"]+=1
                flows[port]["cwnd"].append(flows[port]["pif"])
        
                flows[port]["windows"].append(int(flows[port]["bif"]))
            line_count+=1
            total_bytes+=int(packet.get("frame.len"))
            #print(line_count, total_bytes)
        return flows[port]['times'], flows[port]['windows']



def get_rtt_list(filename, SRC, DST, length = None):
    df_data = load_tcpdump_pcap_csv(filename, length)
    delay_list = []
    current_time_list = []

    record_need_ack_dict = {}

    last_ack = 0

    for index, row in df_data.iterrows():
        time_relative = row['frame.time_relative']
        seq = row['tcp.seq']
        ack = row['tcp.ack']
        tcp_len = row['tcp.len']
    
        if ack == 1:
            if row['ip.dst'] == DST:
                need_ack = seq + tcp_len
                record_need_ack_dict[need_ack] = time_relative
        else:
            if row['ip.dst'] != DST:
                if ack - last_ack >= 0:
                    for key in list(record_need_ack_dict.keys()):
                        if key < ack:
                            del record_need_ack_dict[key]
                        elif key == ack:
                            rtt = time_relative - record_need_ack_dict[key]
                            delay_list.append(rtt)
                            current_time_list.append(time_relative)
                            del record_need_ack_dict[key]

                    last_ack = ack
                else:
                    print('tmp debug', ack - last_ack, ack, last_ack)
        
    return current_time_list, delay_list


def get_packets_in_flight_old2(filename, length = None, interval=0.01):
    df_data = load_tcpdump_pcap_csv(filename, length)
    inflight_packets = 0
    inflight_packets_list = []
    current_time_list = []
    record_need_ack_dict = {}

    interval_cnt = interval
    for index, row in df_data.iterrows():
        if row['ip.dst'] == DST:    # mm -> hombre
            seq = row['tcp.seq']
            time_relative = row['frame.time_relative']
            frame.len = row['frame.len']
            record_need_ack_dict[seq] = frame.len
            inflight_packets += frame.len

            if time_relative >= interval_cnt:
                inflight_packets_list.append(inflight_packets)
                current_time_list.append(time_relative)
                interval_cnt += interval
        else:    # hombre -> mm
            ack = row['tcp.ack']
            for key in list(record_need_ack_dict.keys()):
                if record_need_ack_dict[key] <= ack:
                    inflight_packets -= record_need_ack_dict[key]
                    del record_need_ack_dict[key]
            time_relative = row['frame.time_relative']
            if time_relative >= interval_cnt:
                current_time_list.append(time_relative)
                inflight_packets_list.append(inflight_packets)
                interval_cnt += interval
    return current_time_list, inflight_packets_list


def get_packets_in_flight_old(filename, length = None):
    df_data = load_tcpdump_pcap_csv(filename, length)
    inflight_packets = 0
    inflight_packets_list = []
    current_time_list = []

    record_ack_list = []
    last_time = 0
    for index, row in df_data.iterrows():
        if row['ip.dst'] == DST:    # mm -> hombre
            # Only caculate ACK flag and ACK-PSH flag 
            if row['tcp.flags'] == '0x00000010' or row['tcp.flags'] == '0x00000018':
                seq = row['tcp.seq']
                time_relative = row['frame.time_relative']

                if len(record_ack_list) > 0:
                    while time_relative > record_ack_list[0][0] and record_ack_list[0][0] > last_time:
                        inflight_packets -= record_ack_list[0][1]
                        inflight_packets_list.append(inflight_packets)
                        current_time_list.append(record_ack_list[0][0])
                        if len(record_ack_list) > 0:
                            del record_ack_list[0]
                        if len(record_ack_list) == 0:
                            break

                potential_ack_df = df_data[
                    (df_data['tcp.ack'] == seq)
                    & (df_data['ip.dst'] != DST)
                    & (df_data['tcp.flags'] == '0x00000010')
                    & (df_data['frame.time_relative'] > time_relative)
                ]
                # Filter every packets without ack. (usually every first [P.] packet.)
                if not potential_ack_df.empty:
                    snd_size = row['frame.len']
                    inflight_packets += snd_size
                    inflight_packets_list.append(inflight_packets)
                    current_time_list.append(time_relative)

                    # Only select latest ack packet.
                    potential_ack = potential_ack_df.iloc[0]
                    potential_ack_time = potential_ack['frame.time_relative']
                    record_ack_list.append((potential_ack_time, snd_size))
                    record_ack_list = sorted(record_ack_list, key = lambda x: x[0])
                last_time = time_relative
    return current_time_list, inflight_packets_list


def get_delay_list_old(filename, length = None):
    df_data = load_tcpdump_pcap_csv(filename, length)
    delay_list = []
    current_time_list = []
    for index, row in df_data.iterrows():
        if row['ip.dst'] == DST:    # mm -> hombre
            # Only caculate ACK flag and ACK-PSH flag 
            if row['tcp.flags'] == '0x00000010' or row['tcp.flags'] == '0x00000018':
                seq = row['tcp.seq']
                if seq > 500:
                    time_relative = row['frame.time_relative']
                    potential_ack_df = df_data[
                        (df_data['tcp.ack'] == seq)
                        & (df_data['ip.dst'] != DST)
                        & (df_data['tcp.flags'] == '0x00000010')
                        & (df_data['frame.time_relative'] > time_relative)
                    ]
                    # Filter every packets without ack. (usually every first [P.] packet.)
                    if not potential_ack_df.empty:
                        # Only select latest ack packet.
                        potential_ack = potential_ack_df.iloc[0]
                        potential_ack_time = potential_ack['frame.time_relative']
                        delay_list.append(potential_ack_time - time_relative)
                        current_time_list.append(time_relative)
    return current_time_list, delay_list


def get_throughput_list(filename, SRC, DST, length = None, interval = 0.1):
    df_data = load_tcpdump_pcap_csv(filename, length)
    snd_data = df_data[df_data['ip.dst'] == DST]
    max_time = math.ceil(snd_data.iloc[-1]['frame.time_relative'])
    min_time = math.ceil(snd_data.iloc[0]['frame.time_relative'])

    time_list = []
    throughput_list = []

    width = 0.02

    for i in np.arange(min_time, max_time, interval):
        throughput = snd_data[(snd_data['frame.time_relative'] > i - width) & (snd_data['frame.time_relative'] < i)]['frame.len'].sum()
        time_list.append(i)
        # Bytes/s -> Mbits/s
        throughput_list.append(throughput / width / 1000000 * 8)
    return time_list, throughput_list


def get_loss_list_old(filename, SRC, DST, expected_max_rtt = 0.1):
    df_data = load_tcpdump_pcap_csv(filename)
    snd_data = df_data[df_data['ip.dst'] == DST]
    rcv_data = df_data[df_data['ip.dst'] != DST]

    loss_time_list = []

    for index, row in snd_data.iterrows():
        seq = row['tcp.seq']
        tcp_len = row['tcp.len']
        relative_time = row['frame.time_relative']

        if seq > 1:
            need_ack = seq + tcp_len
            ack_df = rcv_data[(rcv_data['tcp.ack'] == need_ack) & (rcv_data['frame.time_relative'] > relative_time)]
            if len(ack_df) > 0: 
                time_diff = ack_df.iloc[0]['frame.time_relative'] - relative_time
                if time_diff > expected_max_rtt:
                    # packet loss
                    loss_time_list.append(relative_time)
            else:
                # packet loss
                loss_time_list.append(relative_time)

    return loss_time_list


def get_loss_list(filename, DST, expected_max_rtt = 0):
    df_data = load_tcpdump_pcap_csv(filename)
    snd_data = df_data[df_data['ip.dst'] == DST]

    loss_time_list = []

    for index, row in snd_data.iterrows():
        seq = row['tcp.seq']
        tcp_len = row['tcp.len']
        relative_time = row['frame.time_relative']

        if seq > 1:
            # repeat_df = snd_data[(snd_data['tcp.seq'] == seq) & (snd_data['tcp.len'] == tcp_len)]
            repeat_df = snd_data[(snd_data['tcp.seq'] == seq)]
            if len(repeat_df) > 1:
                for _, repeat_row in repeat_df.iterrows():
                    time_diff = repeat_row['frame.time_relative'] - relative_time
                    if time_diff > expected_max_rtt:
                        # packet loss
                        loss_time_list.append(relative_time)

    return loss_time_list


def get_loss_list_from_two_tcpdump_old(filename_a, filename_b, SRC, DST):
    df_data_a = load_tcpdump_pcap_csv(filename_a)
    snd_data_a = df_data_a[df_data_a['ip.dst'] == DST]
    df_data_b = load_tcpdump_pcap_csv(filename_b)
    snd_data_b = df_data_b[df_data_b['ip.dst'] == DST]

    loss_time_list = []
    # allowed_delta_time = 1
    for index, row in snd_data_a.iterrows():
        seq = row['tcp.seq']
        relative_time = row['frame.time_relative']

        if seq > 1:
            # check_b = snd_data_b[(snd_data_b['tcp.seq'] == seq) & (snd_data_b['frame.time_relative'] < relative_time + allowed_delta_time) & (snd_data_b['frame.time_relative'] > relative_time - allowed_delta_time)]
            check_b = snd_data_b[snd_data_b['tcp.seq'] == seq]
            if len(check_b) < 1:
                loss_time_list.append(relative_time)

    return loss_time_list


def get_loss_list_from_two_tcpdump(filename_a, filename_b, SRC, DST):
    # 读取a.csv和b.csv中的数据
    with open(filename_a, 'r') as a_file, open(filename_b, 'r') as b_file:
        a_reader = csv.reader(a_file)
        b_reader = csv.reader(b_file)
        # seqa_list = [row for row in a_reader]
        # seqb_list = [row for row in b_reader]
        # 从a.csv中读取所有符合过滤条件的行
        seqa_list = [row for row in a_reader if row[6] == DST]
        # 从b.csv中读取所有符合过滤条件的行
        seqb_list = [row for row in b_reader if row[6] == DST]

    # 初始化指针pointera和pointerb
    pointera = 0
    pointerb = 0

    # 创建一个空的结果列表
    c_list = []

    init_time = float(seqa_list[0][10])

    # 比较seqa_list和seqb_list中的元素
    while pointera < len(seqa_list) and pointerb < len(seqb_list):
        if float(seqa_list[pointera][3]) == float(seqb_list[pointerb][3]):
            # 如果两个元素的tcp.seq相等，将pointera和pointerb都向后移动一位
            pointera += 1
            pointerb += 1
        elif float(seqa_list[pointera][3]) < float(seqb_list[pointerb][3]):
            # 如果pointera指向的元素的tcp.seq小于pointerb指向的元素的tcp.seq，
            # 如果pointera后一个元素的tcp.seq等于pointerb指向的元素的tcp.seq，则将pointera向后移动一位
            if pointera > 0 and float(seqa_list[pointera ][3]) == float(seqb_list[pointerb+1][3]):
                pointera += 1
                if pointera > 0 and float(seqa_list[pointera ][3]) == float(seqb_list[pointerb+2][3]):
                    pointera += 1
            else:
                # 否则将pointera指向的整行数据添加到结果列表中，并将pointera向后移动一位
                c_list.append(float(seqa_list[pointera][10]) - init_time)
                pointera += 1
        else:
            pointerb += 1

    # 如果seqa_list中还有剩余的元素
    while pointera < len(seqa_list):
        c_list.append(float(seqa_list[pointera][10]) - init_time)
        pointera += 1
    return c_list



if __name__ == '__main__':
    # current_time_list, inflight_packets_list = get_packets_in_flight('cubic_md5_100Mb_lat10ms_uplinkloss1%.pcap.csv')
    # current_time_list, delay_list = get_delay_list('cubic_md5_100Mb_lat10ms_uplinkloss1%.pcap.csv')
    # current_time_list, throughput_list = get_throughput_list('cubic_md5_100Mb_lat10ms_uplinkloss1%.pcap.csv', None, 0.1)
    # print(current_time_list, throughput_list)
    # loss_time_list = get_loss_list('replay_wireless_cubic.pcap.csv', 0)
    # loss_time_list = get_loss_list('loss_edith.pcap.csv', 0.001)
    # loss_time_list = get_loss_list_from_two_tcpdump('loss_ga_eth0.pcap.csv', 'loss_ga_eth1.pcap.csv')
    loss_time_list = get_loss_list_from_two_tcpdump('loss_Ga_eth0_mm1_cubic.pcap.csv', 'loss_Ga_eth1_mm1_cubic.pcap.csv')
    print(loss_time_list)
    print(len(loss_time_list))


